#
# To launch the pipeline:
# 1. Ensure you have correctly-formatted file named configure_me.py
#
# 2. > python runme.py --sequences <filepath1> --ortholist <filepath2>
#
# where, <filepath1> points to a FASTA-formatted text file containing all the sequences (unaligned) for all the species and orthologs in this study.  This file will probably be very large.
# and where <filepath2> points to a text file defining the orthology between FASTA sequences in <filepath1>.  Each line of <filepath2> should contain a whitespace-separated list of sequence names (using names in <filepath1>).  For example, the following line:
#       aaa1 aaa2 aab4 aac8
# declares that sequences aaa1, aaa2, aab4 and aac8 are orthologous, and it is therefore appropriate to align these sequences and infer their shared evolutionary history.
#
#
# Output:
# - a directory named ERG_GENES, containing unaligned FASTA files for each orthology group.
# - a directory named MSAS, containing both FASTA and PHYLIP alignments for all the groups in ERG_GENES.  MSAS will also contain all PHYML output for each ML inference session on each PHYLIP file.

import re, os
from configure_me import *
from argParser import *
from dendropy import Tree

argp = ArgParser(sys.argv)

# A path to a file with all the sequences
# Each sequence name should be formatted as follows: genename_[speciesname],
# where genename can be arbitrary for every sequence, but speciesname should be the same
# across sequences corresponding to the same species.
SEQUENCE_COLLECTION_PATH = argp.getArg("--sequences")

# A filepath where each line contains a list of sequence names that correspond to orthologous genes.  
# The sequences are assumed to all exist in the file specified by --sequences.
# The first ortholog on each line will be the name of the line's set.
# For example, if a line contains "Cytb_S.cerevisae Cytb_C.albicans Cytb_K.lactis",
# then the name of the set will be "Cytb_S.cerevisae."
ORTHOLIST_PATH = argp.getArg("--ortholist")

USE_MPI = argp.getOptionalArg("--mpi")

GENE_DIR = "ERG_GENES"
MSA_DIR = "MSAS"

# returns [sequences, species],
# where sequences = {}, key = sequence name, value = sequence
# and species = {}, key = species name, value = array of sequence names
def parse_all_sequences():
    sequences = {} 
    species = {} 
    
    fin = open(SEQUENCE_COLLECTION_PATH, "r")    
    last_seq_name = ""
    for line in fin.xreadlines():
        if line.startswith(">"):
            last_seq_name = line.strip()[1:]
        else:
            sequences[last_seq_name] = line.strip()
            # Now grab the species name for this sequence:
            speciesname = last_seq_name.split("_[")[1:]
            speciesname = re.sub("\]", "", speciesname[0])
            if species.__contains__(speciesname):
                species[speciesname].append( last_seq_name )
            else:
                species[speciesname] = [last_seq_name]
    fin.close()

    print "\n. I found " + sequences.__len__().__str__() + " sequences."
    print "\n. I found " + species.__len__().__str__() + " species:"   
    for sp in species:
        print sp + " " + species[sp].__len__().__str__() + " genes"
    
    return [sequences, species]

# returns a list of phylip files for the orthologs which are "easy"---
# that is, they are represented in all the species.
def find_easy_genes(sequences, species):
    print "\n. Finding the set of genes that exist in all target species..."

    fin = open(ORTHOLIST_PATH, "r")
    fout = open("easy_genes.txt", "w")
    count_easy = 0
    count_genes = 0
    easy_genes = []
    for line in fin.xreadlines():
        species_in_this_gene = []
        count_genes += 1
        tokens = line.split()
        genename = tokens[0].split("_[")[0]
        genename = genename.split("_")[1:][0]
        for t in tokens:
            speciesname = t.split("_[")[1:]
            speciesname = re.sub("\]", "", speciesname[0])
            species_in_this_gene.append( speciesname )
        easy = True
        for s in species: # ensure that every species is represented in this gene tree
            if False == species_in_this_gene.__contains__( s ):
                easy = False
        if easy and species_in_this_gene.__len__() == species.__len__():
            fout.write( genename + "\n")
            count_easy += 1
            easy_genes.append(genename + ".phylip")
        print genename, easy
    fin.close()
    fout.close()
    print (count_easy / (count_genes * 1.0)), "percent of genes are found in all target species."
    return easy_genes

# returns a list of gene files  
def write_gene_files(sequences):
    gene_files = []
    if False == os.path.exists(GENE_DIR):
        os.system("mkdir " + GENE_DIR)

    fin = open(ORTHOLIST_PATH, "r")
    for line in fin.xreadlines():
        tokens = line.split()
        genename = tokens[0].split("_[")[0]
        genename = genename.split("_")[1:][0]
        print genename + " contains " + tokens.__len__().__str__() + " genes."
        # We write a fasta file containing only sequences from this gene
        fout = open(GENE_DIR + "/" + genename + ".fasta", "w")
        gene_files.append(GENE_DIR + "/" + genename + ".fasta")
        for t in tokens:
            outline = ">" + t + "\n" + sequences[t] + "\n"
            fout.write(outline)
        fout.close()
    fin.close()
    return gene_files

def align_gene_files(gene_files):
    msa_files = []
    if False == os.path.exists(MSA_DIR):
        os.system("mkdir " + MSA_DIR)
    if gene_files.__len__() == 0:
        if os.path.exists(GENE_DIR):
            gene_files = os.listdir(GENE_DIR)
    commands = []
    for f in gene_files:
        filenameseed = f.split("/")[1:][0]
        commands.append( MUSCLE + " -i " + f + " -o " + MSA_DIR + "/" + filenameseed)
    
    commandsfile = "runall_align.sh"
    fout = open(commandsfile, "w")
    for c in commands:
        fout.write( c + "\n" )
    fout.close()
    
    if USE_MPI:
        os.system( MPIDISPATCH + " " + commandsfile)
    else:
        os.system( "source " + commandsfile)

    for f in gene_files:
        if os.path.exists(MSA_DIR + "/" + filenameseed):
            msa_files.append(MSA_DIR + "/" + filenameseed)        
    return msa_files

#
# Convert Muscle's FASTA output into PHYLIP alignments
#
def convert_msas_to_phylip(msa_files):
    print "\n. Converting all FASTA to PHYLIP..."

    phy_files = []
    if msa_files.__len__() == 0:
        if os.path.exists(MSA_DIR):
            msa_files = os.listdir(MSA_DIR)
    
    commands = []
    for f in msa_files:
        if f.__contains__("fasta"):
            commands.append( SEQCONVERTER + " -d" + MSA_DIR + "/" + f + " -if -ope -ri" )
            if (f.__contains__("/")):
                f = f.split("/")[1:][0]
            f.replace("fasta", "phy")
            phy_files.append(f)

    commandsfile = "runall_convert.sh"
    fout = open(commandsfile, "w")
    for c in commands:
        fout.write( c + "\n")
    fout.close()
    if USE_MPI:
        os.system( MPIDISPATCH + " " + commandsfile)
    else:
        os.system( "source " + commandsfile)   

    return phy_files

# convert all genename_[speciesname] into just speciesname
# This will facilitate tree comparison later.
def scrub_phy_sequence_names(filenames):
    print "\n. Normalizing sequence names..."
    here = os.getcwd()
    os.chdir(MSA_DIR)
    for f in filenames:
        fin = open(f, "r")
        lines = fin.readlines()
        fin.close()
        fout = open(f, "w")
        fout.write(lines[0])
        for l in lines[1:]:
            tokens = l.split()
            taxa = tokens[0]
            seq = tokens[1]
            taxa = taxa.split("_[")[1:]
            taxa = re.sub("\]", "", taxa[0])
            fout.write(taxa + "\t" + seq + "\n")
        fout.close()
    os.chdir(here)

def infer_trees(phy_files):
    here = os.getcwd()
    os.chdir(MSA_DIR)
    commands = []
    for f in phy_files:
        command = PHYML + " " + PHYML_FLAGS + " --input " + f
        commands.append( command )
 
    commandsfile = "runall_phyml.sh"
    fout = open(commandsfile, "w")
    for c in commands:
        fout.write( c + "\n" )
    fout.close()
    
    #if USE_MPI:
    #    os.system( MPIDISPATCH + " " + commandsfile)
    #else:
    #    os.system( "source " + commandsfile)   
        
    os.chdir(here)
 
# returns a list of genes, a subset of gene_list, that have verified trees.
def verify_trees(gene_list):
    verified_genes = []
    for g in gene_list:
        treefile = MSA_DIR + "/" + g + ".phylip_phyml_tree.txt"
        if False == os.path.exists(treefile):
            print "ERROR: I cannot find the ML tree for gene", g
        else:
            try:
                t = Tree()
                t.read_from_path(treefile, "newick")
            except IOError:
                print "ERROR: Something is wrong with the ML tree for gene", g
        print g, "ML tree was inferred  OK."
        verified_genes.append(g)
    return verified_genes


# input: pointsfile, where each line contains "genex geney d(x,y)\n"
def draw_heatmap():    
    # parse the points file:
    fin = open("compare_trees.txt", "r")
    lines = fin.readlines()
    fin.close()
    points = {} # an array of arrays, points[genex][geney] = d(genex's topo, geney's topo)
    for line in lines:
        line = line.strip()
        tokens = line.split()
        x = tokens[0]
        y = tokens[1]
        dist = tokens[2]
        if points.keys().__contains__(x):
            points[x][y] = dist
        else:
            points[x] = {}
            points[x][y] = dist

    cranscript = "compare_trees.cran"

    # a CSV file
    fout = open("compare_trees.table", "w")
    # write the header line
    ypoints = points[ points.keys()[0] ].keys()
    ypoints.sort()
    line = ','.join( ypoints )
    fout.write("X," + line + "\n")
    xpoints = points.keys()
    xpoints.sort()
    for x in xpoints:
        print x
        line_points = []
        for y in ypoints:
            line_points.append(points[x][y].__str__() )
        line = ','.join( line_points )
        fout.write(x.__str__() + "," + line + "\n")
    fout.close()
    
    # build a CRAN script
    fout = open(cranscript, "w")
    fout.write("t <- read.csv(\"compare_trees.table\", sep=\",\")\n")
    fout.write("row.names(t) <- t$X\n")
    fout.write("t <- t[,2:" + points[ points.keys()[0] ].keys().__len__().__str__() + "]\n")
    fout.write("tm <- data.matrix(t)\n")
    fout.write("pdf(\"compare_trees.heatmap.pdf\")\n")
    # note: xlab and ylab are switched in the following statement because R likes to rotate plots in a confusing way.
    fout.write("h <- heatmap(tm, Rowv=NA, Colv=NA, col=heat.colors(2056), scale=\"column\", margins=c(5,10), ylab=\"" + xlab + "\", xlab=\"" + ylab + "\" )\n")
    fout.write("dev.off()\n")
    fout.close()
    os.system(CRAN + " --save < " + cranscript)


def compare_trees(gene_list):
    print "\n. Comparing all ML gene trees..."

    all_trees = {} # key = gene name, value = dendropy Tree object
    tree_ranks = {} # key = dendropy tree object, value = array of gene names with this same topology 

    # First, move all trees into memory:
    for g in gene_list:
        print g
        treefile = MSA_DIR + "/" + g + ".phylip_phyml_tree.txt"
        tg = Tree()
        tg.read_from_path(treefile, "newick")
        all_trees[g] = tg
        
        # Second, have we seen this topology before?
        # (increment the tree_ranks structure)
        if tree_ranks.__len__() == 0:
            tree_ranks[tg] = [g]
        else:
            for t in tree_ranks:
                if t.symmetric_difference(tg) == 0:
                    tree_ranks[t].append(g)

    # Third, print tree_ranks
    fout = open("tree_ranks.txt", "w")
    for t in tree_ranks:
        line = t.__str__() + "\t" + tree_ranks[t].__len__().__str__() + "\n"
        fout.write(line)
    fout.close()

    # Fourth, compare all topologies and print the comparisons
    fout = open("compare_trees.txt", "w")
    for gx in all_trees:
        for gy in all_trees:
            line = gx + " " + gy + " " +  all_trees[gx].symmetric_difference( all_trees[gy] ).__str__()
            print line
            fout.write(line + "\n")
    fout.close()

def make_consensus_tree(gene_list):
    all_trees_path = MSA_DIR + "/all_trees.trees" 
    fout = open(all_trees_path, "w")
    treefilelist = ""
    for g in gene_list:
        treefile = MSA_DIR + "/" + g + ".phylip_phyml_tree.txt"
        fin = open(treefile, "r")
        fout.write(fin.readlines()[0])
        fin.close()
    fout.close()

    command = "python /common/DendroPy-3.2.0/scripts/sumtrees/sumtrees.py -m8 "
    command += all_trees_path
    command += " > summarize_trees.txt"
    os.system(command)
   
######################################
#
# main()....
#
sequences = []
species = []
gene_files = [] 
msa_files = [] 
phy_files = []

# 1.
#[sequences, species] = parse_all_sequences()

# 2.
#gene_files = write_gene_files( sequences )

# 3.
#msa_files = align_gene_files( gene_files )

# 4.
#phy_files = convert_msas_to_phylip( msa_files )

#phy_files = find_easy_genes(sequences, species)

# 5.
#scrub_phy_sequence_names(phy_files)

# 6.
#infer_trees( phy_files )

# 7.
fin = open("easy_genes.txt", "r")
lines = fin.readlines()
easy_genes = []
for l in lines:
    easy_genes.append(l.strip())
fin.close()
verified_genes = verify_trees(easy_genes)

# 8. compare trees
compare_trees( verified_genes )
#make_consensus_tree( verified_genes )

# 9. plot
draw_heatmap()
